import json
import pandas as pd
import numpy as np
import random
import tensorflow as tf
import os

from megnet.data.crystal import CrystalGraph
from megnet.data.graph import GaussianDistance
from megnet.models import MEGNetModel

from pymatgen.core.structure import Structure
import sys
import logging

GPU_device = "0"
n_fold = 5
seed = 123
GPU_seed = 11111
ep = 5000
callback = tf.keras.callbacks.EarlyStopping(monitor="val_loss", patience=10, restore_best_weights=True)

random.seed(seed)
np.random.seed(seed)
tf.random.set_seed(GPU_seed)
commit_id = str(os.popen('git --no-pager log -1 --oneline --pretty=format:"%h"').read())
dump_model_name = '{0}_{1}_{2}'.format(commit_id, GPU_device, seed)
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = GPU_device
root_logger = logging.getLogger()
for h in root_logger.handlers[:]:
    root_logger.removeHandler(h)
logging.getLogger().setLevel(logging.INFO)
logging.basicConfig(filename=dump_model_name+".log",
        format='%(asctime)s-%(pathname)s[line:%(lineno)d]-%(levelname)s: %(message)s',
        level=logging.INFO)

structures = []
targets = []

test_structures = []
test_targets = []

def load_data(train_idxs, test_idxs, noise_level):
    logging.info("Current train idxs list: {0}; test idxs list: {1}; noise_level: {2}".format(train_idxs, test_idxs, noise_level))
    denoiser_name = "f934e70_none_0_4.hdf5"
    logging.info("Denoiser name is {0}".format(denoiser_name))
    denoiser = MEGNetModel.from_file(denoiser_name)
    cur_idx = 0
    for i in range(n_fold):
        csv_name = './data/' + str(i) + '_cif.csv'
        df = pd.read_csv(csv_name)
        r = list(range(len(df)))
        if i in test_idxs:
            for j in r:
                tmp = Structure.from_str(df['pbe_structure'][j], fmt='cif')
                test_structures.append(tmp)
                test_targets.append(df['pbe_gap'][j])
        elif i in train_idxs:
            for j in r:
                tmp = Structure.from_str(df['pbe_structure'][j], fmt='cif')
                model_output = denoiser.predict_structure(tmp).ravel()
                structures.append(tmp)
                val = max(0, df['pbe_gap'][j]*(1-np.random.normal(noise_level[cur_idx], noise_level[cur_idx], 1)[0]))
                # val = df['pbe_gap'][j]*(1-noise_level[cur_idx])
                if abs(val-model_output) > 0.35:
                    val = model_output
                targets.append(val)
            cur_idx+=1
    pass

def prediction(model, structures, targets):
    MAE = 0
    test_size = len(structures)
    for i in range(test_size):
        model_output = model.predict_structure(structures[i]).ravel()
        err = abs(model_output - targets[i])
        MAE += err
    MAE /= test_size
    return MAE


def train_and_test(load_name, dump_name):
    logging.info("Data for train: {0}, test: {1}.".format(len(structures), len(test_structures)))
    if load_name == "new":
        cur_model = MEGNetModel(nfeat_edge=10, nfeat_global=2, graph_converter=CrystalGraph(bond_converter=GaussianDistance(np.linspace(0, 5, 10), 0.5)))
    else:
        cur_model = MEGNetModel.from_file(load_name)
    l = len(structures)
    s = structures
    t = targets
    c = list(zip(s, t))
    random.shuffle(c)
    s, t = zip(*c)
    ###### train ############
    try:
        cur_model.train(s[:int(0.8*l)], 
                    t[:int(0.8*l)],
                    validation_structures=s[int(0.8*l):],
                    validation_targets=t[int(0.8*l):],
                    callbacks=[callback],
                    save_checkpoint=False,
                    automatic_correction=False,
                    batch_size = 256,
                    epochs=ep)
    except TypeError:
        logging.info('MAE of {tag} is: {mae}'.format(tag="all_together", mae='nan'))
    else:
        mae = prediction(cur_model, test_structures, test_targets)
        logging.info('Ordered structures MAE of {tag} is: {mae}'.format(tag="all_together", mae=mae))
    cur_model.save_model(dump_name)

def clear_all_lst():
    global structures, targets, test_structures, test_targets
    structures = []
    targets = []
    test_structures = []
    test_targets = []

def main():
    clear_all_lst()
    idxs = [0, 1, 2, 3]
    i = "none"
    load_data(train_idxs=[0], test_idxs=[4], noise_level=[0.3])
    last_name = "new"
    new_name = "{commitid}_{foldid}_{gpu}_1.hdf5".format(commitid=commit_id, foldid=i, gpu=GPU_device)
    train_and_test(load_name=last_name, dump_name=new_name)
    last_name = new_name

    clear_all_lst()
    load_data(train_idxs=[1], test_idxs=[4], noise_level=[0.2])
    new_name = "{commitid}_{foldid}_{gpu}_2.hdf5".format(commitid=commit_id, foldid=i, gpu=GPU_device)
    train_and_test(load_name=last_name, dump_name=new_name)
    last_name = new_name
    
    clear_all_lst()
    load_data(train_idxs=[2], test_idxs=[4], noise_level=[0.1])
    new_name = "{commitid}_{foldid}_{gpu}_3.hdf5".format(commitid=commit_id, foldid=i, gpu=GPU_device)
    train_and_test(load_name=last_name, dump_name=new_name)
    last_name = new_name
 
    clear_all_lst()
    load_data(train_idxs=[3], test_idxs=[4], noise_level=[0.0])
    new_name = "{commitid}_{foldid}_{gpu}_4.hdf5".format(commitid=commit_id, foldid=i, gpu=GPU_device)
    train_and_test(load_name=last_name, dump_name=new_name)
    last_name = new_name


if __name__ == "__main__":
    main()
